import torch
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from torch import nn


class CNNSVM(nn.Module):
    def __init__(self,device=None):
        super().__init__()
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # 假设我们将一个CNN特征提取器和SVM分类器组合起来
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(1, 10, kernel_size=5),
            nn.ReLU(),
            nn.Conv2d(10, 20, kernel_size=5),
            nn.ReLU()
        )
        self.clf = make_pipeline(StandardScaler(), SVC(gamma='auto'))

    def forward(self, x, return_features=False):
        x = self.feature_extractor(x)
        if return_features:
            return x.view(x.size(0), -1)  # Return CNN output features
        # Further steps to process features with an SVM would go here
        return x

